# -*- coding: UTF-8 -*-
from math import log
from decision_tree.util.FileUtil import FileUtil

"""
    决策树计算工具
"""


class MathUtil:
    """
    计算数据集D的经验熵H(D)
    @:param data_set 数据集D
    """

    @staticmethod
    def calcExperienceEnt(dataSet):
        # print("计算经验熵，数据集如下:")
        # print(dataSet)
        numEntries = len(dataSet)
        # 数据集分类字典
        labelCounts = {}
        for featVec in dataSet:
            currentLabel = featVec[-1]
            if currentLabel not in labelCounts.keys():
                labelCounts[currentLabel] = 0
            labelCounts[currentLabel] += 1
        # 计算经验熵
        shannon_ent = 0.0
        for key in labelCounts:
            prob = float(labelCounts[key]) / numEntries
            # 以2为底求对数
            shannon_ent -= prob * log(prob, 2)
        # print("经验熵为：" + str(shannon_ent))
        return shannon_ent

    """
    计算数据集D的条件经验熵H(D|A)
    @:param data_set    数据集D
    @:param feat_index  特征下标，因小于数据集最后一列下标
    """

    @staticmethod
    def calcConditionalExperienceEnt(dataSet, featIndex):
        featList = [example[featIndex] for example in dataSet]
        uniqueVals = set(featList)
        newEntropy = 0.0
        for value in uniqueVals:
            subDataSet = MathUtil.splitDataSet(dataSet, featIndex, value)
            prob = len(subDataSet) / float(len(dataSet))
            newEntropy += prob * MathUtil.calcExperienceEnt(subDataSet)
        return newEntropy

    """
    计算数据集D特征为A的信息增益
    @:param data_set    数据集D
    @:param feat_index  特征下标，因小于数据集最后一列下标
    """

    @staticmethod
    def calcInformationGain(dataSet, featIndex):
        return MathUtil.calcExperienceEnt(dataSet) - MathUtil.calcConditionalExperienceEnt(dataSet, featIndex)


    @staticmethod
    def splitDataSet(dataSet, axis, value):
        retDataSet = []
        for featVec in dataSet:
            if featVec[axis] == value:
                # reducedFeatVec = featVec[:axis]
                # reducedFeatVec.extend(featVec[axis + 1:])
                # retDataSet.append(reducedFeatVec)
                retDataSet.append(featVec)
        return retDataSet


# dateSet, labels = FileUtil.readFile('../test/test2.txt')
# print(MathUtil.calcExperienceEnt(dateSet))
# print("年龄-条件经验熵")
# print(MathUtil.calcConditionalExperienceEnt(dateSet, 0))
# print("年龄-信息增益")
# print(MathUtil.calcInformationGain(dateSet, 0))

